use webrtc::rtp::{packet::Packet, header::Header};
use webrtc_util::Marshal;
use std::sync::atomic::AtomicBool;
use std::sync::{Arc, Mutex};
use tokio::net::UdpSocket;
use webrtc::track::track_remote::TrackRemote;
use std::fs::{read_to_string, write};
use std::path::Path;
use base64::{engine::general_purpose::STANDARD, Engine};
use crate::get_path::get_path;
use once_cell::sync::Lazy;
use std::sync::atomic::Ordering;

// 控制捕获信息不重复打印
pub static STATUS: Lazy<Arc<AtomicBool>> = Lazy::new(|| Arc::new(AtomicBool::new(false)));

/// 构造单一 NALU 的 RTP 包
fn build_nalu_packet(
    nalu: &[u8],
    ssrc: u32,
    seq: u16,
    timestamp: u32,
    payload_type: u8,
    marker: bool,
) -> Packet {
    Packet {
        header: Header {
            version: 2,
            payload_type,
            sequence_number: seq,
            timestamp,
            ssrc,
            marker,
            ..Default::default()
        },
        payload: nalu.to_vec().into(),
    }
}

/// 判断是否是关键帧（IDR）的起始包，支持单一 NAL 和 FU-A 分片格式
fn is_idr_start(payload: &[u8]) -> bool {
    if payload.is_empty() {
        return false;
    }

    let nal_type = payload[0] & 0x1F;

    match nal_type {
        5 => true, // 单一 NAL 单元，IDR
        28 => {
            if payload.len() < 2 {
                return false;
            }
            let fu_header = payload[1];
            let start_bit = fu_header & 0x80 != 0;
            let nal_unit_type = fu_header & 0x1F;
            start_bit && nal_unit_type == 5
        }
        _ => false,
    }
}

/// 更新 SDP 文件中的 SPS/PPS 行
fn update_sdp_file(sps: &[u8], pps: &[u8]) {
    let sps_b64 = STANDARD.encode(sps);
    let pps_b64 = STANDARD.encode(pps);
    let fmtp_line = format!(
        "a=fmtp:102 packetization-mode=1; sprop-parameter-sets={},{}",
        sps_b64, pps_b64
    );

    let path = get_path().join("sdp").join("VideoSDP.sdp");
    let original = match read_to_string(&path) {
        Ok(content) => content,
        Err(_) => {
            println!("⚠️ SDP 文件不存在，跳过写入");
            return;
        }
    };

    let updated = if original.contains("a=fmtp:102") {
        original
            .lines()
            .map(|line| {
                if line.starts_with("a=fmtp:102") {
                    fmtp_line.clone()
                } else {
                    line.to_string()
                }
            })
            .collect::<Vec<_>>()
            .join("\n")
    } else {
        let mut lines = Vec::new();
        let mut inserted = false;
        for line in original.lines() {
            lines.push(line.to_string());
            if !inserted && line.trim() == "a=rtpmap:102 H264/90000" {
                lines.push(fmtp_line.clone());
                inserted = true;
            }
        }
        lines.join("\n")
    };

    if let Err(e) = write(path, updated) {
        println!("❌ 写入 SDP 文件失败: {}", e);
    } else {
        println!("✅ 已更新 SDP 文件中的 SPS/PPS");
    }
}

pub async fn send_rtp_to_udp(
    track: Arc<TrackRemote>,
    target_addr: &str,
    payload_type: u8,
    is_video: bool,
) {
    let socket = UdpSocket::bind("0.0.0.0:0").await.unwrap();

    let ssrc = 0x12345678;
    let mut sequence_number: u16 = 1;

    let sps_cache = Arc::new(Mutex::new(None::<Vec<u8>>));
    let pps_cache = Arc::new(Mutex::new(None::<Vec<u8>>));
    let sdp_written = Arc::new(Mutex::new(false));

    loop {
        match track.read_rtp().await {
            Ok(packet) => {
                let rtp_packet = packet.0;
                let payload = &rtp_packet.payload;

                if payload.is_empty() {
                    continue;
                }

                if is_video {
                    let nal_type = payload[0] & 0x1F;

                    // STAP-A 聚合包
                    if nal_type == 24 {
                        let mut offset = 1;
                        while offset + 2 <= payload.len() {
                            let nalu_size = u16::from_be_bytes([payload[offset], payload[offset + 1]]) as usize;
                            offset += 2;
                            if offset + nalu_size > payload.len() {
                                break;
                            }
                            let nalu_data = &payload[offset..offset + nalu_size];
                            let sub_nal_type = nalu_data[0] & 0x1F;
                            match sub_nal_type {
                                7 => {
                                    let mut sps_lock = sps_cache.lock().unwrap();
                                    *sps_lock = Some(nalu_data.to_vec());
                                    if !STATUS.load(Ordering::SeqCst) {
                                        println!("✅ 从 STAP-A 捕获 SPS");
                                    }
                                    STATUS.store(true, Ordering::SeqCst);
                                }
                                8 => {
                                    let mut pps_lock = pps_cache.lock().unwrap();
                                    *pps_lock = Some(nalu_data.to_vec());
                                    if !STATUS.load(Ordering::SeqCst) {
                                        println!("✅ 从 STAP-A 捕获 PPS");
                                    }
                                    STATUS.store(true, Ordering::SeqCst);
                                }
                                _ => {}
                            }
                            offset += nalu_size;
                        }
                    } else {
                        match nal_type {
                            7 => {
                                let mut sps_lock = sps_cache.lock().unwrap();
                                *sps_lock = Some(payload.to_vec());
                                if !STATUS.load(Ordering::SeqCst) {
                                    println!("✅ 捕获 SPS");
                                }
                                STATUS.store(true, Ordering::SeqCst);
                            }
                            8 => {
                                let mut pps_lock = pps_cache.lock().unwrap();
                                *pps_lock = Some(payload.to_vec());
                                if !STATUS.load(Ordering::SeqCst) {
                                    println!("✅ 捕获 PPS");
                                }
                                STATUS.store(true, Ordering::SeqCst);
                            }
                            _ => {}
                        }
                    }

                    // 写入 SDP 文件（只写一次）
                    let sps_data = sps_cache.lock().unwrap().clone();
                    let pps_data = pps_cache.lock().unwrap().clone();

                    if let (Some(sps), Some(pps)) = (sps_data.clone(), pps_data.clone()) {
                        let should_write = {
                            let mut written_flag = sdp_written.lock().unwrap();
                            if !*written_flag {
                                *written_flag = true;
                                true
                            } else {
                                false
                            }
                        };

                        if should_write {
                            update_sdp_file(&sps, &pps);
                        }
                    }
                }

                let mut forwarded_packet = rtp_packet.clone();
                forwarded_packet.header.sequence_number = sequence_number;
                forwarded_packet.header.ssrc = ssrc;

                let raw_forwarded_packet = forwarded_packet.marshal().unwrap();
                socket.send_to(&raw_forwarded_packet, target_addr).await.unwrap();
                sequence_number = sequence_number.wrapping_add(1);
            }
            Err(err) => {
                eprintln!("❌ RTP 读取失败: {:?}", err);
                break;
            }
        }
    }
}